# -*- coding: utf-8 -*-

# 导入必要的库
from keras.models import load_model  # 加载预训练模型
from keras import backend as K       # Keras后端
import numpy as np                   # 数值计算库
import librosa                       # 音频处理库
from python_speech_features import mfcc  # 提取MFCC特征
import pickle                        # 序列化工具
import glob                          # 文件路径匹配
import tensorflow as tf

# 获取所有.wav音频文件路径
wavs = glob.glob('data/*.wav')

# 加载字符映射字典和MFCC特征统计信息  
with open(r'./dictionary.pkl', 'rb') as fr:
    [char2id, id2char, mfcc_mean, mfcc_std] = pickle.load(fr)

# 设置MFCC特征维度
mfcc_dim = 13

# 加载预训练的语音识别模型
model = load_model('asr.h5')

# 随机选择一个音频文件
index = np.random.randint(len(wavs))
print(wavs[index])

# 加载音频文件
audio, sr = librosa.load(wavs[index])

# 计算音频能量并去除静音段
# energy = librosa.feature.rmse(audio)  # 计算均方根能量
energy = librosa.feature.rms(y=audio)
frames = np.nonzero(energy >= np.max(energy) / 5)  # 找到能量大于最大能量1/5的帧
indices = librosa.core.frames_to_samples(frames)[1]  # 将帧转换为采样点
# 截取非静音部分
audio = audio[indices[0]:indices[-1]] if indices.size else audio[0:0]

# 提取MFCC特征并进行标准化
X_data = mfcc(audio, sr, numcep=mfcc_dim, nfft=551)
X_data = (X_data - mfcc_mean) / (mfcc_std + 1e-14)  # 标准化处理
print(X_data.shape)

# 读取对应的文本标签
with open(wavs[index] + '.trn', 'r', encoding='utf8') as fr:
    label = fr.readlines()[0]
    print(label)

# 使用模型进行预测
pred = model.predict(np.expand_dims(X_data, axis=0))  # 扩展维度以匹配模型输入

# 使用TensorFlow的CTC解码器解码预测结果
decoded, _ = tf.nn.ctc_greedy_decoder(
    tf.transpose(pred, perm=[1, 0, 2]), 
    [X_data.shape[0]]
)
pred_ids = tf.sparse.to_dense(decoded[0]).numpy().flatten().tolist()

# 将预测的ID序列转换为字符并打印结果
print(''.join([id2char[i] for i in pred_ids if i in id2char]))